import os
from pydub import AudioSegment

def parse_time(time_str):
    """将时间字符串转换为秒数"""
    parts = list(map(int, time_str.split(':')))
    if len(parts) == 2:
        return parts[0] * 60 + parts[1]
    elif len(parts) == 3:
        return parts[0] * 3600 + parts[1] * 60 + parts[2]
    else:
        raise ValueError("时间格式无效，请使用 mm:ss 或 hh:mm:ss 格式")

def main():
    print("音频处理工具 - 支持时间截取与格式转换")
    
    # 获取音频文件路径
    while True:
        file_path = input("请输入音频文件路径：").strip()
        if os.path.exists(file_path):
            break
        print("文件不存在，请重新输入。")

    # 加载音频文件
    try:
        audio = AudioSegment.from_file(file_path)
    except Exception as e:
        print(f"无法读取音频文件：{e}")
        return

    total_duration = audio.duration_seconds  # 总时长（秒）
    print(f"音频总时长：{total_duration:.2f} 秒")

    # --- 修改点：先选择操作模式 ---
    print("\n请选择操作：")
    print("1. 截取指定时间段")
    print("2. 转换整个文件为WAV")
    while True:
        choice = input("请输入选项（1或2）：").strip()
        if choice in ['1', '2']:
            break
        print("无效选项，请输入1或2。")

    # --- 修改点：根据选择执行不同逻辑 ---
    if choice == '2':
        # 转换整个文件为WAV
        output_path = os.path.splitext(file_path)[0] + "_full.wav"
        audio.export(output_path, format="wav")
        print(f"\n✅ 已将整个音频转换为WAV格式：{output_path}")
    else:
        # 用户选择了截取指定时间段，现在才询问时间
        # 获取开始时间
        while True:
            start_time = input("请输入开始时间（如 01:12 或 01:23:45）：").strip()
            try:
                start_seconds = parse_time(start_time)
                if start_seconds < 0 or start_seconds >= total_duration:
                    print("开始时间超出音频范围！")
                    continue
                break
            except ValueError as e:
                print(e)

        # 获取结束时间
        while True:
            end_time = input("请输入结束时间（如 01:30 或 01:23:45）：").strip()
            try:
                end_seconds = parse_time(end_time)
                if end_seconds <= start_seconds:
                    print("结束时间必须在开始时间之后！")
                elif end_seconds > total_duration:
                    print("结束时间超出音频范围！")
                else:
                    break
            except ValueError as e:
                print(e)

        # 执行截取操作
        start_ms = start_seconds * 1000
        end_ms = end_seconds * 1000
        selected_audio = audio[start_ms:end_ms]
        
        # 生成输出文件名
        original_name = os.path.splitext(file_path)[0]
        safe_start = start_time.replace(":", "_")
        safe_end = end_time.replace(":", "_")
        output_path = f"{original_name}_from_{safe_start}_to_{safe_end}.wav"
        
        selected_audio.export(output_path, format="wav")
        print(f"\n✅ 已截取 {start_time}-{end_time} 段落并保存为：{output_path}")

if __name__ == "__main__":
    main()